from typing import Dict
from rainforeLearn.gomoku.v2.environment.env import GomokuEnvironment
from rainforeLearn.gomoku.v2.agents.dqn_agent import GomokuDQNAgent
from rainforeLearn.gomoku.v2.configs.config import GomokuDQNConfig
from rainforeLearn.gomoku.v2.train.constants.train_constants import TrainConstants


class ModelEvaluator:
    """模型评估器 - 负责模型性能评估"""
    
    def __init__(self, env: GomokuEnvironment, agent: GomokuDQNAgent, config: GomokuDQNConfig):
        self.env = env
        self.agent = agent
        self.config = config
    
    def evaluate(self, num_games: int) -> float:
        """评估性能"""
        print(f"🎯 开始评估 ({num_games} 局)...")
        
        wins = 0
        total_games = 0
        
        for _ in range(num_games):
            if self._play_single_evaluation_game():
                wins += 1
            total_games += 1
        
        win_rate = wins / total_games if total_games > 0 else 0
        print(f"📊 评估结果: 胜率 {win_rate:.3f} ({wins}/{total_games})")
        
        return win_rate
    
    def _play_single_evaluation_game(self) -> bool:
        """进行单局评估游戏，返回是否获胜"""
        state = self.env.reset()
        game_length = 0
        
        while not self.env.game_over:
            action = self._select_evaluation_action(state)
            
            if action == TrainConstants.INVALID_ACTION:
                break
            
            state, _, _, _ = self.env.step(action)
            game_length += 1
            
            # 防止无限循环
            if game_length > self.config.max_game_length:
                break
        
        return self.env.winner == TrainConstants.AI_PLAYER
    
    def _select_evaluation_action(self, state: Dict) -> int:
        """选择评估动作"""
        if self.env.current_player == TrainConstants.AI_PLAYER:
            return self.agent.select_action(state, epsilon=0)
        else:
            return self.agent.select_action(state, epsilon=0.0)